# config.py
# 集中管理所有联邦学习实验的超参数和配置。

import torch
import numpy as np
import random

def get_config(dataset_name='CIFAR10', scheme='scheme1'):
    """
    根据数据集名称和方案返回相应的配置。
    参数:
        dataset_name (str): 数据集名称, 目前支持 'CIFAR10'。
        scheme (str): 实验方案, 支持 'scheme1', 'scheme2', 'scheme3'。
    返回:
        dict: 包含实验所需全部配置的字典。
    """
    # --- 通用配置 ---
    # 种子现在由主函数管理，确保同一次运行中所有scheme使用相同seed，不同次运行使用不同seed
    
    config = {
        "project_name": "federated_label_shift",
        "device": "cuda:0" if torch.cuda.is_available() else "cpu",  # 优先使用GPU
        "seed": 42,  # 默认种子，会被主函数覆盖
        "num_workers": 2,

        # --- 联邦学习配置 ---
        "num_clients": 3,
        
        # --- 数据集配置 ---
        "dataset": dataset_name,
        "data_path": "./data",
        
        # --- 客户端数据异构性配置 ---
        # 为3个客户端分别指定类别比例的狄利克雷分布alpha值
        # 调整为更温和的差异：alpha值在1-4之间，确保有差异但不极端
        "clients_dirichlet_alphas": [0.5, 0.5, 0.5],
        # 每个客户端的数据总量
        "client_data_sizes": [8000, 9000, 10000],
        # 数据抽样方式：True=允许重叠抽样, False=不重叠抽样
        "allow_data_overlap": True,

        # --- 服务器目标域配置 ---
        # 注意：现在服务器数据来自测试集的均匀分割，不再使用Dirichlet分布
        # 保留这些参数以保持向后兼容，但实际不再使用
        "server_target_dirichlet_alpha": 0.3,  # 不再使用
        "server_unlabeled_data_size": 10000,   # 不再使用，现在固定为测试集的一半
        
        # --- 模型配置 ---
        # 使用CLIP的预训练视觉模型
        "base_model_name": "clip_vit_b_32",

        # --- 训练超参数 (仅用于在CLIP之上训练分类头) ---
        "batch_size": 128,
        "classifier_epochs": 5, # 训练分类头的轮数
        "classifier_lr": 1e-3,

        # --- 校准与估计算法超参数 ---
        "calibrator_max_iter": 150,
        "em_max_iter": 1000,
        "em_tol": 1e-6,
        
        # --- 实验方案配置 ---
        "scheme": scheme,
    }

    if dataset_name == 'CIFAR10':
        config.update({
            "num_classes": 10,
            "class_names": ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
        })
    # 在这里为其他数据集添加接口
    # elif dataset_name == 'AnotherDataset':
    #     config.update({ ... })
    else:
        raise ValueError(f"Unsupported dataset: {dataset_name}")
        
    # 确保配置长度匹配客户端数量
    assert len(config["clients_dirichlet_alphas"]) == config["num_clients"], "Alpha list must match num_clients."
    assert len(config["client_data_sizes"]) == config["num_clients"], "Data size list must match num_clients."
        
    return config
